import os
import numpy as np
import tensorflow as tf
# import input_data
import lenet5.lenet5_model as model
import lenet5.readdata as rd

N_CLASSES = 2
IMG_W = 208  # resize the image, if the input image is too large, training will be very slow.
IMG_H = 208
BATCH_SIZE = 64
CAPACITY = 2000
MAX_STEP = 10000  # with current parameters, it is suggested to use MAX_STEP>10k
learning_rate = 0.0001  # with current parameters, it is suggested to use learning rate<0.0001


# %%
def run_training1():
    # you need to change the directories to yours.
    #    train_dir = '/home/hjxu/PycharmProjects/01_cats_vs_dogs/data/train/'
    logs_train_dir = 'logs\\recordstrain\\'
    #
    #    train, train_label = input_data.get_files(train_dir)
    # 第一步获取数据，batch和lable
    tfrecords_file = 'catvsdogtrain.tfrecords'
    train_batch, train_label_batch = rd.read_and_decode(tfrecords_file, batch_size=BATCH_SIZE)
    train_batch = tf.cast(train_batch, dtype=tf.float32)
    train_label_batch = tf.cast(train_label_batch, dtype=tf.int64)
    # 第二步得到模型处理后的softmax_linear，
    train_logits = model.inference(train_batch, BATCH_SIZE, N_CLASSES)
    # 计算损失函数
    train_loss = model.losses(train_logits, train_label_batch)
    # 将损失函数和学习率代入训练
    train_op = model.trainning(train_loss, learning_rate)
    # 计算，验证
    train__acc = model.evaluation(train_logits, train_label_batch)
    # tf.summary都是tensorboard相关
    # summary_op = tf.summary.merge_all()
    sess = tf.Session()
    # train_writer = tf.summary.FileWriter(logs_train_dir, sess.graph)
    # 定义tf.train保存对象
    saver = tf.train.Saver()

    sess.run(tf.global_variables_initializer())
    # 申明一个tf.train.Coordinator类来协同多线程
    coord = tf.train.Coordinator()
    threads = tf.train.start_queue_runners(sess=sess, coord=coord)

    try:
        for step in np.arange(MAX_STEP):
            if coord.should_stop():
                break
                # 开始跑训练步骤
            _, tra_loss, tra_acc = sess.run([train_op, train_loss, train__acc])
            # 每50步计算一次，并且保存、可视化到tensorboard
            if step % 50 == 0:
                print('Step %d, train loss = %.2f, train accuracy = %.2f%%' % (step, tra_loss, tra_acc * 100.0))
                # summary_str = sess.run(summary_op)
                # train_writer.add_summary(summary_str, step)
            # 训练完成之后，保存模型
            if step % 2000 == 0 or (step + 1) == MAX_STEP:
                checkpoint_path = os.path.join(logs_train_dir, 'model.ckpt')
                saver.save(sess, checkpoint_path, global_step=step)

    except tf.errors.OutOfRangeError:
        print('Done training -- epoch limit reached')
    finally:
        coord.request_stop()

    coord.join(threads)
    sess.close()


# %% Evaluate one image
# when training, comment the following codes.


# from PIL import Image
# import matplotlib.pyplot as plt
#
# def get_one_image(train):
#    '''Randomly pick one image from training data
#    Return: ndarray
#    '''
#    n = len(train)
#    ind = np.random.randint(0, n)
#    img_dir = train[ind]
#
#    image = Image.open(img_dir)
#    plt.imshow(image)
#    image = image.resize([208, 208])
#    image = np.array(image)
#    return image
#
# def evaluate_one_image():
#    '''Test one image against the saved models and parameters
#    '''
#
#    # you need to change the directories to yours.
#    train_dir = '/home/kevin/tensorflow/cats_vs_dogs/data/train/'
#    train, train_label = input_data.get_files(train_dir)
#    image_array = get_one_image(train)
#
#    with tf.Graph().as_default():
#        BATCH_SIZE = 1
#        N_CLASSES = 2
#
#        image = tf.cast(image_array, tf.float32)
#        image = tf.image.per_image_standardization(image)
#        image = tf.reshape(image, [1, 208, 208, 3])
#        logit = model.inference(image, BATCH_SIZE, N_CLASSES)
#
#        logit = tf.nn.softmax(logit)
#
#        x = tf.placeholder(tf.float32, shape=[208, 208, 3])
#
#        # you need to change the directories to yours.
#        logs_train_dir = '/home/kevin/tensorflow/cats_vs_dogs/logs/train/'
#
#        saver = tf.train.Saver()
#
#        with tf.Session() as sess:
#
#            print("Reading checkpoints...")
#            ckpt = tf.train.get_checkpoint_state(logs_train_dir)
#            if ckpt and ckpt.model_checkpoint_path:
#                global_step = ckpt.model_checkpoint_path.split('/')[-1].split('-')[-1]
#                saver.restore(sess, ckpt.model_checkpoint_path)
#                print('Loading success, global_step is %s' % global_step)
#            else:
#                print('No checkpoint file found')
#
#            prediction = sess.run(logit, feed_dict={x: image_array})
#            max_index = np.argmax(prediction)
#            if max_index==0:
#                print('This is a cat with possibility %.6f' %prediction[:, 0])
#            else:
#                print('This is a dog with possibility %.6f' %prediction[:, 1])


# %%


if __name__ == '__main__':
    run_training1()
